{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "37861d7e",
   "metadata": {},
   "source": [
    "# Fourier Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb35109c",
   "metadata": {},
   "source": [
    "We analyze feature maps in Fourier space to demonstrate that ***MSA is a low-pass filter*** as shown in Fig. 2, Fig. 8, and Fig. D.2. Fourier analysis shows that MSAs reduce high-frequency signals, while Convs (and MLPs) amplify them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "515971e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "# check whether run in Colab\n",
    "root = \".\"\n",
    "if \"google.colab\" in sys.modules:\n",
    "    print(\"Running in Colab.\")\n",
    "    !pip3 install matplotlib\n",
    "    !pip3 install einops==0.4.1\n",
    "    !pip3 install timm==0.5.4\n",
    "    !git clone https://github.com/xxxnell/how-do-vits-work.git\n",
    "    root = \"./how-do-vits-work\"\n",
    "    sys.path.append(root)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad0723a5",
   "metadata": {},
   "source": [
    "## Load and Prepare Pretrained Models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dbd1eab",
   "metadata": {},
   "source": [
    "Divide the pretrained `timm` models into the sequences of blocks (the cell below provides the snippets for ResNet-50 and ViT-Ti):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "008f2bf4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import timm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "name = \"resnet50\"\n",
    "model = timm.create_model(name, pretrained=True)\n",
    "\n",
    "# model → blocks. `blocks` is a sequence of blocks\n",
    "blocks = [\n",
    "    nn.Sequential(model.conv1, model.bn1, model.act1, model.maxpool),\n",
    "    *model.layer1,\n",
    "    *model.layer2,\n",
    "    *model.layer3,\n",
    "    *model.layer4,\n",
    "    nn.Sequential(model.global_pool, model.fc)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dece74f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "import timm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "name = \"vit_tiny_patch16_224\"\n",
    "model = timm.create_model(name, pretrained=True)\n",
    "\n",
    "\n",
    "class PatchEmbed(nn.Module):\n",
    "    def __init__(self, model):\n",
    "        super().__init__()\n",
    "        self.model = copy.deepcopy(model)\n",
    "        \n",
    "    def forward(self, x, **kwargs):\n",
    "        x = self.model.patch_embed(x)\n",
    "        cls_token = self.model.cls_token.expand(x.shape[0], -1, -1)\n",
    "        x = torch.cat((cls_token, x), dim=1)\n",
    "        x = self.model.pos_drop(x + self.model.pos_embed)\n",
    "        return x\n",
    "\n",
    "\n",
    "class Residual(nn.Module):\n",
    "    def __init__(self, *fn):\n",
    "        super().__init__()\n",
    "        self.fn = nn.Sequential(*fn)\n",
    "        \n",
    "    def forward(self, x, **kwargs):\n",
    "        return self.fn(x, **kwargs) + x\n",
    "    \n",
    "    \n",
    "class Lambda(nn.Module):\n",
    "    def __init__(self, fn):\n",
    "        super().__init__()\n",
    "        self.fn = fn\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.fn(x)\n",
    "\n",
    "\n",
    "def flatten(xs_list):\n",
    "    return [x for xs in xs_list for x in xs]\n",
    "\n",
    "\n",
    "# `blocks` is a sequence of blocks\n",
    "blocks = [\n",
    "    PatchEmbed(model),\n",
    "    *flatten([[Residual(b.norm1, b.attn), Residual(b.norm2, b.mlp)] \n",
    "              for b in model.blocks]),\n",
    "    nn.Sequential(model.norm, Lambda(lambda x: x[:, 0]), model.head),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec79651c",
   "metadata": {},
   "source": [
    "In this experiment, we use pretrained models for ImageNet-1K, not for CIFAR-100. The image size of the CIFAR-100 is too small for this experiment."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fed203fb",
   "metadata": {},
   "source": [
    "## Plot Fourier Transformed Feature Maps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3316201a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This cell build off https://github.com/facebookresearch/mae\n",
    "import requests\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "from PIL import Image\n",
    "from einops import rearrange, reduce, repeat\n",
    "\n",
    "\n",
    "imagenet_mean = np.array([0.485, 0.456, 0.406])\n",
    "imagenet_std = np.array([0.229, 0.224, 0.225])\n",
    "\n",
    "# load a sample ImageNet-1K image -- use the full val dataset for precise results\n",
    "xs = [\n",
    "    \"https://user-images.githubusercontent.com/930317/158025258-e9a5a454-99de-4d22-bc93-b217cdf06abb.jpeg\",\n",
    "]\n",
    "xs = [Image.open(requests.get(x, stream=True).raw) for x in xs]\n",
    "xs = [x.resize((224, 224)) for x in xs]\n",
    "xs = [np.array(x) / 255. for x in xs]\n",
    "xs = np.stack(xs)\n",
    "\n",
    "assert xs.shape[1:] == (224, 224, 3)\n",
    "\n",
    "# normalize by ImageNet mean and std\n",
    "xs = xs - imagenet_mean\n",
    "xs = xs / imagenet_std\n",
    "xs = rearrange(torch.tensor(xs, dtype=torch.float32), \"b h w c -> b c h w\")\n",
    "\n",
    "# accumulate `latents` by collecting hidden states of a model\n",
    "latents = []\n",
    "with torch.no_grad():\n",
    "    for block in blocks:\n",
    "        xs = block(xs)\n",
    "        latents.append(xs)\n",
    "        \n",
    "if name in [\"vit_tiny_patch16_224\", \"pit_ti_224\"]:  # for ViT: Drop CLS token\n",
    "    latents = [latent[:,1:] for latent in latents]\n",
    "latents = latents[:-1]  # drop logit (output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7350a20",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "import matplotlib.cm as cm\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.collections import LineCollection\n",
    "\n",
    "\n",
    "def fourier(x):  # 2D Fourier transform\n",
    "    f = torch.fft.fft2(x)\n",
    "    f = f.abs() + 1e-6\n",
    "    f = f.log()\n",
    "    return f\n",
    "\n",
    "\n",
    "def shift(x):  # shift Fourier transformed feature map\n",
    "    b, c, h, w = x.shape\n",
    "    return torch.roll(x, shifts=(int(h/2), int(w/2)), dims=(2, 3))\n",
    "\n",
    "\n",
    "def make_segments(x, y):  # make segment for `plot_segment`\n",
    "    points = np.array([x, y]).T.reshape(-1, 1, 2)\n",
    "    segments = np.concatenate([points[:-1], points[1:]], axis=1)\n",
    "    return segments\n",
    "\n",
    "\n",
    "def plot_segment(ax, xs, ys, cmap_name=\"plasma\"):  # plot with cmap segments\n",
    "    z = np.linspace(0.0, 1.0, len(ys))\n",
    "    z = np.asarray(z)\n",
    "    \n",
    "    cmap = cm.get_cmap(cmap_name)\n",
    "    norm = plt.Normalize(0.0, 1.0)\n",
    "    segments = make_segments(xs, ys)\n",
    "    lc = LineCollection(segments, array=z, cmap=cmap_name, norm=norm,\n",
    "                        linewidth=2.5, alpha=1.0)\n",
    "    ax.add_collection(lc)\n",
    "\n",
    "    colors = [cmap(x) for x in xs]\n",
    "    ax.scatter(xs, ys, color=colors, marker=marker, zorder=100)\n",
    "    \n",
    "\n",
    "# Fourier transform feature maps\n",
    "fourier_latents = []\n",
    "for latent in latents:  # `latents` is a list of hidden feature maps in latent spaces\n",
    "    latent = latent.cpu()\n",
    "    \n",
    "    if len(latent.shape) == 3:  # for ViT\n",
    "        b, n, c = latent.shape\n",
    "        h, w = int(math.sqrt(n)), int(math.sqrt(n))\n",
    "        latent = rearrange(latent, \"b (h w) c -> b c h w\", h=h, w=w)\n",
    "    elif len(latent.shape) == 4:  # for CNN\n",
    "        b, c, h, w = latent.shape\n",
    "    else:\n",
    "        raise Exception(\"shape: %s\" % str(latent.shape))\n",
    "    latent = fourier(latent)\n",
    "    latent = shift(latent).mean(dim=(0, 1))\n",
    "    latent = latent.diag()[int(h/2):]  # only use the half-diagonal components\n",
    "    latent = latent - latent[0]  # visualize 'relative' log amplitudes \n",
    "                                 # (i.e., low-freq amp - high freq amp)\n",
    "    fourier_latents.append(latent)\n",
    "    \n",
    "\n",
    "# A. Plot Fig 2a: \"Relative log amplitudes of Fourier transformed feature maps\"\n",
    "fig, ax1 = plt.subplots(1, 1, figsize=(3.3, 4), dpi=150)\n",
    "for i, latent in enumerate(reversed(fourier_latents[:-1])):\n",
    "    freq = np.linspace(0, 1, len(latent))\n",
    "    ax1.plot(freq, latent, color=cm.plasma_r(i / len(fourier_latents)))\n",
    "    \n",
    "ax1.set_xlim(left=0, right=1)\n",
    "\n",
    "ax1.set_xlabel(\"Frequency\")\n",
    "ax1.set_ylabel(\"$\\Delta$ Log amplitude\")\n",
    "\n",
    "from matplotlib.ticker import FormatStrFormatter\n",
    "ax1.yaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
    "ax1.xaxis.set_major_formatter(FormatStrFormatter('%.1fπ'))\n",
    "\n",
    "\n",
    "# B. Plot Fig 8: \"Relative log amplitudes of high-frequency feature maps\"\n",
    "if name == \"resnet50\":  # for ResNet-50\n",
    "    pools = [4, 8, 14]\n",
    "    msas = []\n",
    "    marker = \"D\"\n",
    "elif name == \"vit_tiny_patch16_224\":  # for ViT-Ti\n",
    "    pools = []\n",
    "    msas = [1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23,]\n",
    "    marker = \"o\"\n",
    "else:\n",
    "    import warnings\n",
    "    warnings.warn(\"The configuration for %s are not implemented.\" % name, Warning)\n",
    "    pools, msas = [], []\n",
    "    marker = \"s\"\n",
    "\n",
    "depths = range(len(fourier_latents))\n",
    "\n",
    "# Normalize\n",
    "depth = len(depths) - 1\n",
    "depths = (np.array(depths)) / depth\n",
    "pools = (np.array(pools)) / depth\n",
    "msas = (np.array(msas)) / depth\n",
    "\n",
    "fig, ax2 = plt.subplots(1, 1, figsize=(6.5, 4), dpi=120)\n",
    "plot_segment(ax2, depths, [latent[-1] for latent in fourier_latents])  # high-frequency component\n",
    "\n",
    "for pool in pools:\n",
    "    ax2.axvspan(pool - 1.0 / depth, pool + 0.0 / depth, color=\"tab:blue\", alpha=0.15, lw=0)\n",
    "for msa in msas:\n",
    "    ax2.axvspan(msa - 1.0 / depth, msa + 0.0 / depth, color=\"tab:gray\", alpha=0.15, lw=0)\n",
    "    \n",
    "ax2.set_xlabel(\"Normalized depth\")\n",
    "ax2.set_ylabel(\"$\\Delta$ Log amplitude\")\n",
    "ax2.set_xlim(0.0, 1.0)\n",
    "\n",
    "from matplotlib.ticker import FormatStrFormatter\n",
    "ax2.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ad0e127",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
